#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：tiny-llm 
@File    ：model.py
@IDE     ：PyCharm 
@Author  ：XJ
@Date    ：2025/3/20 16:54 
'''
import torch
import torch.nn as nn
from transformers import GPT2Tokenizer


class TinyLLM(nn.Module):
    def __init__(self, vocab_size=50257, embed_dim=128, num_layers=4):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, embed_dim)
        self.layers = nn.ModuleList([
            nn.TransformerEncoderLayer(d_model=embed_dim, nhead=4)
            for _ in range(num_layers)
        ])
        self.fc = nn.Linear(embed_dim, vocab_size)

    def forward(self, x):
        x = self.embed(x)
        for layer in self.layers:
            x = layer(x)
        return self.fc(x)


# 示例用法
if __name__ == "__main__":
    tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
    model = TinyLLM()
    inputs = tokenizer("Hello, AI!", return_tensors="pt")["input_ids"]
    outputs = model(inputs)
    print("输出维度:", outputs.shape)  # 输出维度: torch.Size([1, 3, 50257])